#!/usr/bin/env python
# -*- encoding: utf-8 -*-
"""
@author:hengk
@contact: hengk@foxmail.com
@datetime:2020-05-07 14:07
"""
from backbone.mobilenetV3 import MobileNetV3
from torchvision import transforms
from easydict import EasyDict
import yaml
import torch
import cv2
import os

def prehandle(name,longer_size):
    image = cv2.imread(name)
    h, w, c = image.shape
    scale = longer_size / max(h, w)
    image = cv2.resize(image, (0, 0), fx=scale, fy=scale)
    return transforms.ToTensor()(image).unsqueeze(0)
if __name__ == '__main__':
    f = open("config/default.yaml", 'r', encoding='utf-8')
    cfg = yaml.load(f.read(), Loader=yaml.FullLoader)
    cfg = EasyDict(cfg)
    # 定义网络结构
    os.environ["CUDA_VISIBLE_DEVICES"] = cfg.train.gpu
    net = MobileNetV3()
    net.eval()
    #加载模型参数
    checkpoint = torch.load(os.path.join("model","8.pth"))
    net.load_state_dict(checkpoint)



    name1 = os.path.join(cfg.test.data_dir,"1.jpg")
    name2 = os.path.join(cfg.test.data_dir,"2.jpg")
    name3 = os.path.join(cfg.test.data_dir,"3.jpg")



    image1 = prehandle(name1,cfg.test.longer_size)
    image2 = prehandle(name2, cfg.test.longer_size)
    image3 = prehandle(name3, cfg.test.longer_size)



    output1 = net(image1)
    output2 = net(image2)
    output3 = net(image3)

    a = torch.cosine_similarity(output1, output2)
    b = torch.cosine_similarity(output2, output3)
    print(a.item())
    print(b.item())